/**
 * Copyright 2012 Brigham Young University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package edu.byu.nlp.cluster.mom;

import java.util.Arrays;
import java.util.logging.Logger;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;

import edu.byu.nlp.cluster.BasicProbabilisticClustering;
import edu.byu.nlp.cluster.Dataset;
import edu.byu.nlp.cluster.ProbabilisticClusterer;
import edu.byu.nlp.cluster.ProbabilisticClustering;
import edu.byu.nlp.cluster.ProbabilisticModel;
import edu.byu.nlp.cluster.em.ParameterInitializer;
import edu.byu.nlp.data.SparseFeatureVector;
import edu.byu.nlp.math.optimize.old.ConsoleIterationListener;
import edu.byu.nlp.math.optimize.old.DifferentiableFunction;
import edu.byu.nlp.math.optimize.old.LBFGSMinimizer;
import edu.byu.nlp.util.DoubleArrays;
import edu.byu.nlp.util.Pair;

/**
 * @author rah67
 *
 */
public class GradientInductive implements ProbabilisticClusterer {

	private static Logger logger = Logger.getLogger(GradientInductive.class.getName());
	
	private final double alpha;
	private final double beta;
	private final int maxIterations;
	private final ParameterInitializer<MoMParameters> initializer;

	public GradientInductive(double alpha, double beta, int maxIterations,
			ParameterInitializer<MoMParameters> initializer) {
		this.alpha = alpha;
		this.beta = beta;
		this.maxIterations = maxIterations;
		this.initializer = initializer;
	}

	@VisibleForTesting
	ObjectiveFunction newObjectiveFunction(Dataset data, int numClusters) {
		return new ObjectiveFunction(data, numClusters);
	}
	
	@VisibleForTesting
	final class ObjectiveFunction implements DifferentiableFunction {

		private final Dataset data;
		private final int numClusters;
		
		private final LoadingCache<double[], Pair<Double, double[]>> cache;
		
		public ObjectiveFunction(Dataset data, int numClusters) {
			this.data = data;
			this.numClusters = numClusters;
			this.cache = CacheBuilder.newBuilder()
					.maximumSize(1)
					.weakKeys()
					.weakValues()
					.build(new ValueAndDerivative());
		}
		
		/** {@inheritDoc} */
		@Override
		public int dimension() {
			// |\theta| + |\phi| + |\lambda_\theta| + |\lambda_\phi| = L + LF + 1 + L = 2L + LF + 1 = L(2 + F) + 1
			return numClusters * (data.getNumFeatures() + 2) + 1;
		}

		/** {@inheritDoc} */
		@Override
		public double valueAt(double... x) {
			return cache.getUnchecked(x).getFirst();
		}

		/** {@inheritDoc} */
		@Override
		public double[] derivativeAt(double... x) {
			return cache.getUnchecked(x).getSecond();
		}
		
		@VisibleForTesting
		ValueAndDerivative.Model newModel(double[] params) {
			return new ValueAndDerivative().newModel(params);
		}
		
		@VisibleForTesting
		ValueAndDerivative.Gradient newGradient(int size) {
			return new ValueAndDerivative().newGradient(size);
		}
		
		@VisibleForTesting
		final class ValueAndDerivative extends CacheLoader<double[], Pair<Double, double[]>> {
			
			// Used for testing
			private Model newModel(double[] params) {
				return new Model(params);
			}
			
			@VisibleForTesting
			final class Model {
				
				private final double[] params;
				private final double[] logJoint;
				
				private Model(double[] params) {
					this.params = params;
					this.logJoint = new double[numClusters];
				}
				
				/**
				 * This class maintains ownership of array. Clone it if it needs to persist between calls to this method.
				 */
				public double[] logJoint(SparseFeatureVector instance) {
					System.arraycopy(params, 0, logJoint, 0, numClusters);
					instance.preMultiplyAsColumnAndAddTo(params, numClusters, logJoint);
					return logJoint;
				}
				
			}
			
			// Used for testing
			private Gradient newGradient(int size) {
				return new Gradient(size);
			}
			
			@VisibleForTesting
			final class Gradient {
				
				private final double[] gradient;

				// TODO : make a static factory method if this class is ever refactored and moved to a static context
				public Gradient(int size) {
					this(new double[size]);
				}
				
				private Gradient(double[] gradient) {
					this.gradient = gradient;
				}
				
				private int sizeOfPhi() {
					return numClusters * data.getNumFeatures();
				}
				
				private int thetaLagrangeMultiplierPos() {
					return numClusters + sizeOfPhi();
				}
				
				private int phiLagrangeMultiplierPos(int k) {
					return thetaLagrangeMultiplierPos() + k + 1;
				}
				
				public void setTheta(double value) {
					Arrays.fill(gradient, 0, numClusters, value);
				}
				
				public void addToTheta(double[] x) {
					DoubleArrays.addToSelf(gradient, 0, x, 0, numClusters);
				}
				
				public void divideTheta(double[] x) {
					DoubleArrays.pointwiseDivideToSelf(gradient, 0, x, 0, numClusters);
				}
				
				public void setPhi(double value) {
					Arrays.fill(gradient, numClusters, sizeOfPhi(), value);
				}
				
				public void addToPhiWeighted(SparseFeatureVector x, double[] weight) {
					x.preMultiplyAndAddTo(weight, gradient, numClusters);
				}
				
				public void dividePhi(double[] x) {
					DoubleArrays.pointwiseDivideToSelf(gradient, numClusters, x, numClusters, sizeOfPhi());
				}
				
				public void addThetaLagrangeMultiplier(double[] x) {
					DoubleArrays.addToRangeToSelf(gradient, 0, numClusters, x[thetaLagrangeMultiplierPos()]);
				}
				
				public void addPhiLagrangeMultiplier(double[] x) {
					int colPos = numClusters; 
					for (int f = 0; f < data.getNumFeatures(); f++) {
						DoubleArrays.addRangeToSelf(gradient, colPos, x, phiLagrangeMultiplierPos(0), numClusters);
						colPos += numClusters;
					}
				}
				
				public void setThetaLagrangeMultiplierToSumOfThetaMinusOne(double[] x) {
					gradient[thetaLagrangeMultiplierPos()] = DoubleArrays.sumOfRange(x, 0, numClusters) - 1;
				}
				
				public void setPhiLagrangianMultiplierToSumsOfPhisMinusOne(double[] x) {
					for (int k = 0; k < numClusters; k++) {
						gradient[phiLagrangeMultiplierPos(k)] = DoubleArrays.sumOfRange(x, 0, data.getNumFeatures()) - 1;
					}
				}
				
			}
			
			@Override
			public Pair<Double, double[]> load(double[] x) {
				// We need log parameters
				DoubleArrays.logToSelf(x);
				Model model = new Model(x);

				double value = 0.0;
				value += (alpha - 1) * DoubleArrays.sumOfRange(x, 0, numClusters);
				value += (beta - 1) * DoubleArrays.sumOfRange(x, numClusters, numClusters * data.getNumFeatures());

				Gradient gradient = new Gradient(x.length);
				// \partial \theta_k = alpha - 1
				gradient.setTheta(alpha - 1);
				// \partial \phi_kv = beta - 1
				gradient.setPhi(beta - 1);

				// FIXME : might not work with seen data
				for (SparseFeatureVector instance : data.unlabeledData()) {
					// will contain p(y|x)
					double[] pYGivenX = model.logJoint(instance);
					// value += log \sum_k p(y_i, x_i)
					value += DoubleArrays.logNormalizeToSelf(pYGivenX);
					DoubleArrays.expToSelf(pYGivenX);
					
					// \partial \theta_k += p(Y_i = k | x_i; \theta, \phi)
					gradient.addToTheta(pYGivenX);
					// \partial \phi_kv += p(Y_i = k | x_i; \theta, \phi) * x_iv
					gradient.addToPhiWeighted(instance, pYGivenX);
				}
				
				// No longer want log parameters; we need the actual thing
				DoubleArrays.expToSelf(x);
				
				// \partial \theta_k /= \theta_k
				gradient.divideTheta(x);
				// \partial \theta_k += \lambda_\theta
				gradient.addThetaLagrangeMultiplier(x);

				// \partial \phi_kv /= \phi_kv
				gradient.dividePhi(x);
				// \partial \phi_kv += \lambda_\theta_k
				gradient.addPhiLagrangeMultiplier(x);
				
				// \partial \lambda_\theta = \sum_k \theta_k - 1
				gradient.setThetaLagrangeMultiplierToSumOfThetaMinusOne(x);
				// \partial \lambda_\phi_k = \sum_v \phi_kv - 1
				gradient.setPhiLagrangianMultiplierToSumsOfPhisMinusOne(x);
				
				value += x[gradient.thetaLagrangeMultiplierPos()] * (DoubleArrays.sumOfRange(x, 0, numClusters) - 1);
				for (int k = 0; k < numClusters; k++) {
					double sum = 0.0;
					int index = numClusters + k;
					for (int v = 0; v < data.getNumFeatures(); v++) {
						sum += x[index];
						index += numClusters;
					}
					value += x[gradient.phiLagrangeMultiplierPos(k)] * (sum - 1);
				}
				
				// FIXME : do the negative inline
				value = -value;
				DoubleArrays.multiplyToSelf(gradient.gradient, -1);
				
				return Pair.of(value, gradient.gradient);
			}
		}
	}
	
	public ProbabilisticClustering cluster(Dataset data, int numClusters) {
		
		ObjectiveFunction g = new ObjectiveFunction(data, numClusters);
		double[] initial = initialize(data, numClusters, g.dimension());
		
		LBFGSMinimizer minimizer = new LBFGSMinimizer(maxIterations);
		minimizer.addIterationListener(new ConsoleIterationListener());
		double[] parameters = minimizer.minimize(g, initial, 1e-6);
		
		ProbabilisticModel model = newModel(parameters, numClusters, data.getNumFeatures());
		// FIXME : return correct objective value
		return BasicProbabilisticClustering.newClustering(data, numClusters, model, model, Double.NaN);
	}

	private double[] initialize(Dataset data, int numClusters, int dim) {
		// TODO : a little wasteful with memory, but probably not a problem
		MoMParameters params = initializer.initialize(data, numClusters);
		double[] initial = new double[dim];
		
		// We don't want log parameters
		DoubleArrays.expToSelf(params.getLogPOfY());
		System.arraycopy(params.getLogPOfY(), 0, initial, 0, params.getLogPOfY().length);
		
		// Simultaneously exponentiate and convert to column order
		double[][] pOfXGivenY = params.getLogPOfXGivenY();
		for (int k = 0; k < numClusters; k++) {
			int linearIndex = k;
			for (int f = 0; f < pOfXGivenY[k].length; f++) {
				initial[numClusters + linearIndex] = Math.exp(pOfXGivenY[k][f]);
				linearIndex += numClusters;
			}
		}
		
		// We let the lagrange multipliers stay at 0.0
		return initial;
	}
	
	private ProbabilisticModel newModel(double[] parameters, int numClusters, int numFeatures) {
		double[] pOfY = new double[numClusters];
		System.arraycopy(parameters, 0, pOfY, 0, numClusters);

		double[][] pOfXGivenY = new double[numClusters][numFeatures];
		for (int k = 0; k < numClusters; k++) {
			int linearIndex = k;
			for (int f = 0; f < pOfXGivenY[k].length; f++) {
				pOfXGivenY[k][f] = parameters[numClusters + linearIndex];
				linearIndex += numClusters;
			}
		}
		
		// FIXME : remove distribution check
		MoMParameters param = MoMParameters.fromProbabilities(pOfY, pOfXGivenY, false, true);
		
		return MixtureOfMultinomialsModel.newWithParameters(param, false);
	}

}
